{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TVP-VAR, MCMC, and sparse simulation smoothing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "from importlib import reload\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import statsmodels.api as sm\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from scipy.stats import invwishart, invgamma\n",
    "\n",
    "# Get the macro dataset\n",
    "dta = sm.datasets.macrodata.load_pandas().data\n",
    "dta.index = pd.date_range('1959Q1', '2009Q3', freq='QS')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Background\n",
    "\n",
    "Bayesian analysis of linear Gaussian state space models via Markov chain Monte Carlo (MCMC) methods has become both commonplace and relatively straightforward in recent years, due especially to advances in sampling from the joint posterior of the unobserved state vector conditional on the data and model parameters (see especially Carter and Kohn (1994), de Jong and Shephard (1995), and Durbin and Koopman (2002)). This is particularly useful for Gibbs sampling MCMC approaches.\n",
    "\n",
    "While these procedures make use of the forward/backward application of the recursive Kalman filter and smoother, another recent line of research takes a different approach and constructs the posterior joint distribution of the entire vector of states at once - see in particular Chan and Jeliazkov (2009) for an econometric time series treatment and McCausland et al. (2011) for a more general survey. In particular, the posterior mean and precision matrix are constructed explicitly, with the latter a sparse band matrix. Advantage is then taken of efficient algorithms for Cholesky factorization of sparse band matrices; this reduces memory costs and can improve performance. Following McCausland et al. (2011), we refer to this method as the \"Cholesky Factor Algorithm\" (CFA) approach.\n",
    "\n",
    "The CFA-based simulation smoother has some advantages and some drawbacks compared to that based on the more typical Kalman filter and smoother (KFS).\n",
    "\n",
    "**Advantages of CFA**:\n",
    "\n",
    "- Derivation of the joint posterior distribution is relatively straightforward and easy to understand.\n",
    "- In some cases can be both faster and less memory-intensive than the KFS approach\n",
    "    - In the Appendix at the end of this notebook, we briefly discuss the performance of the two simulation smoothers for the TVP-VAR model. In summary: simple tests on a single machine suggest that for the TVP-VAR model, the CFA and KFS implementations in Statsmodels have about the same runtimes, while both implementations are about twice as fast as the replication code, written in Matlab, provided by Chan and Jeliakov (2009).\n",
    "\n",
    "**Drawbacks of CFA**:\n",
    "\n",
    "The main drawback is that this method has not (at least so far) reached the generality of the KFS approach. For example:\n",
    "\n",
    "- It can not be used with models that have reduced-rank error terms in the observation or state equations.\n",
    "    - One implication of this is that the typical state space model trick of including identities in the state equation to accomodate, for example, higher-order lags in autoregressive models is not applicable. These models can still be handled by the CFA approach, but at the cost of requiring a slightly different implementation for each lag that is included.\n",
    "    - As an example, standard ways of representing ARMA and VARMA processes in state space form do include identities in the observation and/or state equations, and so the basic formulas presented in Chan and Jeliazkov (2009) do not apply immediately to these models.\n",
    "- Less flexibility is available in the state initialization / prior."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementation in Statsmodels\n",
    "\n",
    "A CFA simulation smoother along the lines of the basic formulas presented in Chan and Jeliazkov (2009) has been implemented in Statsmodels.\n",
    "\n",
    "**Notes**:\n",
    "\n",
    "- Therefore, the CFA simulation smoother in Statsmodels so-far only supports the case that the state transition is truly a first-order Markov process (i.e. it does not support a p-th order Markov process that has been stacked using identities into a first-order process).\n",
    "- By contrast, the KFS smoother in Statsmodels is fully general any can be used for any state space model, including those with stacked p-th order Markov processes or other identities in the observation and state equations.\n",
    "\n",
    "Either a KFS or the CFA simulation smoothers can be constructed from a state space model using the `simulation_smoother` method. To show the basic idea, we first consider a simple example."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Local level model\n",
    "\n",
    "A local level model decomposes an observed series $y_t$ into a persistent trend $\\mu_t$ and a transitory error component\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "y_t & = \\mu_t + \\varepsilon_t, \\qquad \\varepsilon_t \\sim N(0, \\sigma_\\text{irregular}^2) \\\\\n",
    "\\mu_t & = \\mu_{t-1} + \\eta_t, \\quad ~ \\eta_t \\sim N(0, \\sigma_\\text{level}^2)\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "This model satisfies the requirements of the CFA simulation smoother because both the observation error term $\\varepsilon_t$ and the state innovation term $\\eta_t$ are non-degenerate - that is, their covariance matrices are full rank.\n",
    "\n",
    "We apply this model to inflation, and consider simulating draws from the posterior of the joint state vector. That is, we are interested in sampling from\n",
    "\n",
    "$$p(\\mu^t \\mid y^t, \\sigma_\\text{irregular}^2, \\sigma_\\text{level}^2)$$\n",
    "\n",
    "where we define $\\mu^t \\equiv (\\mu_1, \\dots, \\mu_T)'$ and $y^t \\equiv (y_1, \\dots, y_T)'$.\n",
    "\n",
    "In Statsmodels, the local level model falls into the more general class of \"unobserved components\" models, and can be contructed as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct a local level model for inflation\n",
    "mod = sm.tsa.UnobservedComponents(dta.infl, 'llevel')\n",
    "\n",
    "# Fit the model's parameters (sigma2_varepsilon and sigma2_eta)\n",
    "# via maximum likelihood\n",
    "res = mod.fit()\n",
    "print(res.params)\n",
    "\n",
    "# Create simulation smoother objects\n",
    "sim_kfs = mod.simulation_smoother()              # default method is KFS\n",
    "sim_cfa = mod.simulation_smoother(method='cfa')  # can specify CFA method"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The simulation smoother objects `sim_kfs` and `sim_cfa` have `simulate` methods that perform simulation smoothing. Each time that `simulate` is called, the `simulated_state` attribute will be re-populated with a new simulated draw from the posterior.\n",
    "\n",
    "Below, we construct 20 simulated paths for the trend, using the KFS and CFA approaches, where the simulation is at the maximum likelihood parameter estimates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nsimulations = 20\n",
    "simulated_state_kfs = pd.DataFrame(\n",
    "    np.zeros((mod.nobs, nsimulations)), index=dta.index)\n",
    "simulated_state_cfa = pd.DataFrame(\n",
    "    np.zeros((mod.nobs, nsimulations)), index=dta.index)\n",
    "\n",
    "for i in range(nsimulations):\n",
    "    # Apply KFS simulation smoothing\n",
    "    sim_kfs.simulate()\n",
    "    # Save the KFS simulated state\n",
    "    simulated_state_kfs.iloc[:, i] = sim_kfs.simulated_state[0]\n",
    "\n",
    "    # Apply CFA simulation smoothing\n",
    "    sim_cfa.simulate()\n",
    "    # Save the CFA simulated state\n",
    "    simulated_state_cfa.iloc[:, i] = sim_cfa.simulated_state[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plotting the observed data and the simulations created using each method below, it is not too hard to see that these two methods are doing the same thing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the inflation data along with simulated trends\n",
    "fig, axes = plt.subplots(2, figsize=(15, 6))\n",
    "\n",
    "# Plot data and KFS simulations\n",
    "dta.infl.plot(ax=axes[0], color='k')\n",
    "axes[0].set_title('Simulations based on KFS approach, MLE parameters')\n",
    "simulated_state_kfs.plot(ax=axes[0], color='C0', alpha=0.25, legend=False)\n",
    "\n",
    "# Plot data and CFA simulations\n",
    "dta.infl.plot(ax=axes[1], color='k')\n",
    "axes[1].set_title('Simulations based on CFA approach, MLE parameters')\n",
    "simulated_state_cfa.plot(ax=axes[1], color='C0', alpha=0.25, legend=False)\n",
    "\n",
    "# Add a legend, clean up layout\n",
    "handles, labels = axes[0].get_legend_handles_labels()\n",
    "axes[0].legend(handles[:2], ['Data', 'Simulated state'])\n",
    "fig.tight_layout();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Updating the model's parameters\n",
    "\n",
    "The simulation smoothers are tied to the model instance, here the variable `mod`. Whenever the model instance is updated with new parameters, the simulation smoothers will take those new parameters into account in future calls to the `simulate` method.\n",
    "\n",
    "This is convenient for MCMC algorithms, which repeatedly (a) update the model's parameters, (b) draw a sample of the state vector, and then (c) draw new values for the model's parameters.\n",
    "\n",
    "Here we will change the model to a different parameterization that yields a smoother trend, and show how the simulated values change (for brevity we only show the simulations from the KFS approach, but simulations from the CFA approach would be the same)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(15, 3))\n",
    "\n",
    "# Update the model's parameterization to one that attributes more\n",
    "# variation in inflation to the observation error and so has less\n",
    "# variation in the trend component\n",
    "mod.update([4, 0.05])\n",
    "\n",
    "# Plot simulations\n",
    "for i in range(nsimulations):\n",
    "    sim_kfs.simulate()\n",
    "    ax.plot(dta.index, sim_kfs.simulated_state[0],\n",
    "            color='C0', alpha=0.25, label='Simulated state')\n",
    "\n",
    "# Plot data\n",
    "dta.infl.plot(ax=ax, color='k', label='Data', zorder=-1)\n",
    "    \n",
    "# Add title, legend, clean up layout\n",
    "ax.set_title('Simulations with alternative parameterization yielding a smoother trend')\n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "ax.legend(handles[-2:], labels[-2:])\n",
    "fig.tight_layout();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Application: Bayesian analysis of a TVP-VAR model by MCMC\n",
    "\n",
    "One of the applications that Chan and Jeliazkov (2009) consider is the time-varying parameters vector autoregression (TVP-VAR) model, estimated with Bayesian Gibb sampling (MCMC) methods. They apply this to model the co-movements in four macroeconomic time series:\n",
    "\n",
    "- Real GDP growth\n",
    "- Inflation\n",
    "- Unemployment rate\n",
    "- Short-term interest rates\n",
    "\n",
    "We will replicate their example, using a very similar dataset that is included in Statsmodels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Subset to the four variables of interest\n",
    "y = dta[['realgdp', 'cpi', 'unemp', 'tbilrate']].copy()\n",
    "y.columns = ['gdp', 'inf', 'unemp', 'int']\n",
    "\n",
    "# Convert to real GDP growth and CPI inflation rates\n",
    "y[['gdp', 'inf']] = np.log(y[['gdp', 'inf']]).diff() * 100\n",
    "y = y.iloc[1:]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(15, 5))\n",
    "y.plot(ax=ax)\n",
    "ax.set_title('Evolution of macroeconomic variables included in TVP-VAR exercise');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### TVP-VAR model\n",
    "\n",
    "**Note**: this section is based on Chan and Jeliazkov (2009) section 3.1, which can be consulted for additional details.\n",
    "\n",
    "The usual (time-invariant) VAR(1) model is typically written:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "y_t & = \\mu + \\Phi y_{t-1} + \\varepsilon_t, \\qquad \\varepsilon_t \\sim N(0, H)\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "where $y_t$ is a $p \\times 1$ vector of variables observed at time $t$ and $H$ is a covariance matrix.\n",
    "\n",
    "The TVP-VAR(1) model generalizes this to allow the coefficients to vary over time according. Stacking all the parameters into a vector according to $\\alpha_t = \\text{vec}([\\mu_t : \\Phi_t])$, where $\\text{vec}$ denotes the operation that stacks columns of a matrix into a vector, we model their evolution over time according to:\n",
    "\n",
    "$$\\alpha_{i,t+1} = \\alpha_{i, t} + \\eta_{i,t}, \\qquad \\eta_{i, t} \\sim N(0, \\sigma_i^2)$$\n",
    "\n",
    "In other words, each parameter evolves independently according to a random walk.\n",
    "\n",
    "Note that there are $p$ coefficients in $\\mu_t$ and $p^2$ coefficients in $\\Phi_t$, so the full state vector $\\alpha$ is shaped $p * (p + 1) \\times 1$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Putting the TVP-VAR(1) model into state-space form is relatively straightforward, and in fact we just have to re-write the observation equation into SUR form:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "y_t & = Z_t \\alpha_t + \\varepsilon_t, \\qquad \\varepsilon_t \\sim N(0, H) \\\\\n",
    "\\alpha_{t+1} & = \\alpha_t + \\eta_t, \\qquad \\eta_t \\sim N(0, \\text{diag}(\\{\\sigma_i^2\\}))\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "$$\n",
    "Z_t = \\begin{bmatrix}\n",
    "1 & y_{t-1}' & 0 & \\dots & &  0 \\\\\n",
    "0 & 0 & 1 & y_{t-1}' &  & 0 \\\\\n",
    "\\vdots & & & \\ddots & \\ddots & \\vdots \\\\\n",
    "0 & 0 & 0 & 0 & 1 & y_{t-1}'  \\\\\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "As long as $H$ is full rank and each of the variances $\\sigma_i^2$ is non-zero, the model satisfies the requirements of the CFA simulation smoother.\n",
    "\n",
    "We also need to specify the initialization / prior for the initial state, $\\alpha_1$. Here we will follow Chan and Jeliazkov (2009) in using $\\alpha_1 \\sim N(0, 5 I)$, although we could also model it as diffuse."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Aside from the time-varying coefficients $\\alpha_t$, the other parameters that we will need to estimate are terms in the covariance matrix $H$ and the random walk variances $\\sigma_i^2$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### TVP-VAR model in Statsmodels\n",
    "\n",
    "Constructing this model programatically in Statsmodels is also relatively straightforward, since there are basically four steps:\n",
    "\n",
    "1. Create a new `TVPVAR` class as a subclass of `sm.tsa.statespace.MLEModel`\n",
    "2. Fill in the fixed values of the state space system matrices\n",
    "3. Specify the initialization of $\\alpha_1$\n",
    "4. Create a method for updating the state space system matrices with new values of the covariance matrix $H$ and the random walk variances $\\sigma_i^2$.\n",
    "\n",
    "To do this, first note that the general state space representation used by Statsmodels is:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "y_t & = d_t + Z_t \\alpha_t + \\varepsilon_t, \\qquad \\varepsilon_t \\sim N(0, H_t) \\\\\n",
    "\\alpha_{t+1} & = c_t + T_t \\alpha_t + R_t \\eta_t, \\qquad \\eta_t \\sim N(0, Q_t) \\\\\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "Then the TVP-VAR(1) model implies the following specializations:\n",
    "\n",
    "- The intercept terms are zero, i.e. $c_t = d_t = 0$\n",
    "- The design matrix $Z_t$ is time-varying but its valeus are fixed as described above (i.e. its values contain ones and lags of $y_t$)\n",
    "- The observation covariance matrix is not time-varying, i.e. $H_t = H_{t+1} = H$\n",
    "- The transition matrix is not time-varying and is equal to the identity matrix, i.e. $T_t = T_{t+1} = I$\n",
    "- The selection matrix $R_t$ is not time-varying and is also equal to the identity matrix, i.e. $R_t = R_{t+1} = I$\n",
    "- The state covariance matrix $Q_t$ is not time-varying and is diagonal, i.e. $Q_t = Q_{t+1} = \\text{diag}(\\{\\sigma_i^2\\})$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. Create a new TVPVAR class as a subclass of sm.tsa.statespace.MLEModel\n",
    "class TVPVAR(sm.tsa.statespace.MLEModel):\n",
    "    # Steps 2-3 are best done in the class \"constructor\", i.e. the __init__ method\n",
    "    def __init__(self, y):\n",
    "        # Create a matrix with [y_t' : y_{t-1}'] for t = 2, ..., T\n",
    "        augmented = sm.tsa.lagmat(y, 1, trim='both', original='in', use_pandas=True)\n",
    "        # Separate into y_t and z_t = [1 : y_{t-1}']\n",
    "        p = y.shape[1]\n",
    "        y_t = augmented.iloc[:, :p]\n",
    "        z_t = sm.add_constant(augmented.iloc[:, p:])\n",
    "\n",
    "        # Recall that the length of the state vector is p * (p + 1)\n",
    "        k_states = p * (p + 1)\n",
    "        super().__init__(y_t, exog=z_t, k_states=k_states)\n",
    "\n",
    "        # Note that the state space system matrices default to contain zeros,\n",
    "        # so we don't need to explicitly set c_t = d_t = 0.\n",
    "\n",
    "        # Construct the design matrix Z_t\n",
    "        # Notes:\n",
    "        # -> self.k_endog = p is the dimension of the observed vector\n",
    "        # -> self.k_states = p * (p + 1) is the dimension of the observed vector\n",
    "        # -> self.nobs = T is the number of observations in y_t\n",
    "        self['design'] = np.zeros((self.k_endog, self.k_states, self.nobs))\n",
    "        for i in range(self.k_endog):\n",
    "            start = i * (self.k_endog + 1)\n",
    "            end = start + self.k_endog + 1\n",
    "            self['design', i, start:end, :] = z_t.T\n",
    "\n",
    "        # Construct the transition matrix T = I\n",
    "        self['transition'] = np.eye(k_states)\n",
    "\n",
    "        # Construct the selection matrix R = I\n",
    "        self['selection'] = np.eye(k_states)\n",
    "\n",
    "        # Step 3: Initialize the state vector as alpha_1 ~ N(0, 5I)\n",
    "        self.ssm.initialize('known', stationary_cov=5 * np.eye(self.k_states))\n",
    "\n",
    "    # Step 4. Create a method that we can call to update H and Q\n",
    "    def update_variances(self, obs_cov, state_cov_diag):\n",
    "        self['obs_cov'] = obs_cov\n",
    "        self['state_cov'] = np.diag(state_cov_diag)\n",
    "\n",
    "    # Finally, it can be convenient to define human-readable names for\n",
    "    # each element of the state vector. These will be available in output\n",
    "    @property\n",
    "    def state_names(self):\n",
    "        state_names = np.empty((self.k_endog, self.k_endog + 1), dtype=object)\n",
    "        for i in range(self.k_endog):\n",
    "            endog_name = self.endog_names[i]\n",
    "            state_names[i] = (\n",
    "                ['intercept.%s' % endog_name] +\n",
    "                ['L1.%s->%s' % (other_name, endog_name) for other_name in self.endog_names])\n",
    "        return state_names.ravel().tolist()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above class defined the state space model for any given dataset. Now we need to create a specific instance of it with the dataset that we created earlier containing real GDP growth, inflation, unemployment, and interest rates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an instance of our TVPVAR class with our observed dataset y\n",
    "mod = TVPVAR(y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Preliminary investigation with ad-hoc parameters in H, Q"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In our analysis below, we will need to begin our MCMC iterations with some initial parameterization. Following Chan and Jeliazkov (2009) we will set $H$ to be the sample covariance matrix of our dataset, and we will set $\\sigma_i^2 = 0.01$ for each $i$.\n",
    "\n",
    "Before discussing the MCMC scheme that will allow us to make inferences about the model, first we can consider the output of the model when simply plugging in these initial parameters. To fill in these parameters, we use the `update_variances` method that we defined earlier and then perform Kalman filtering and smoothing conditional on those parameters.\n",
    "\n",
    "**Warning: This exercise is just by way of explanation - we must wait for the output of the MCMC exercise to study the actual implications of the model in a meaningful way**. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_obs_cov = np.cov(y.T)\n",
    "initial_state_cov_diag = [0.01] * mod.k_states\n",
    "\n",
    "# Update H and Q\n",
    "mod.update_variances(initial_obs_cov, initial_state_cov_diag)\n",
    "\n",
    "# Perform Kalman filtering and smoothing\n",
    "# (the [] is just an empty list that in some models might contain\n",
    "# additional parameters. Here, we don't have any additional parameters\n",
    "# so we just pass an empty list)\n",
    "initial_res = mod.smooth([])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `initial_res` variable contains the output of Kalman filtering and smoothing, conditional on those initial parameters. In particular, we may be interested in the \"smoothed states\", which are $E[\\alpha_t \\mid y^t, H, \\{\\sigma_i^2\\}]$.\n",
    "\n",
    "First, lets create a function that graphs the coefficients over time, separated into the equations for equation of the observed variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_coefficients_by_equation(states):\n",
    "    fig, axes = plt.subplots(2, 2, figsize=(15, 8))\n",
    "\n",
    "    # The way we defined Z_t implies that the first 5 elements of the\n",
    "    # state vector correspond to the first variable in y_t, which is GDP growth\n",
    "    ax = axes[0, 0]\n",
    "    states.iloc[:, :5].plot(ax=ax)\n",
    "    ax.set_title('GDP growth')\n",
    "    ax.legend()\n",
    "\n",
    "    # The next 5 elements correspond to inflation\n",
    "    ax = axes[0, 1]\n",
    "    states.iloc[:, 5:10].plot(ax=ax)\n",
    "    ax.set_title('Inflation rate')\n",
    "    ax.legend();\n",
    "\n",
    "    # The next 5 elements correspond to unemployment\n",
    "    ax = axes[1, 0]\n",
    "    states.iloc[:, 10:15].plot(ax=ax)\n",
    "    ax.set_title('Unemployment equation')\n",
    "    ax.legend()\n",
    "\n",
    "    # The last 5 elements correspond to the interest rate\n",
    "    ax = axes[1, 1]\n",
    "    states.iloc[:, 15:20].plot(ax=ax)\n",
    "    ax.set_title('Interest rate equation')\n",
    "    ax.legend();\n",
    "    \n",
    "    return ax\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we are interested in the smoothed states, which are available in the `states.smoothed` attribute out our results object `initial_res`.\n",
    "\n",
    "As the graph below shows, the initial parameterization implies substantial time-variation in some of the coefficients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here, for illustration purposes only, we plot the time-varying\n",
    "# coefficients conditional on an ad-hoc parameterization\n",
    "\n",
    "# Recall that `initial_res` contains the Kalman filtering and smoothing,\n",
    "# and the `states.smoothed` attribute contains the smoothed states\n",
    "plot_coefficients_by_equation(initial_res.states.smoothed);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Bayesian estimation via MCMC\n",
    "\n",
    "We will now implement the Gibbs sampler scheme described in Chan and Jeliazkov (2009), Algorithm 2.\n",
    "\n",
    "\n",
    "We use the following (conditionally conjugate) priors:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "H & \\sim \\mathcal{IW}(\\nu_1^0, S_1^0) \\\\\n",
    "\\sigma_i^2 & \\sim \\mathcal{IG} \\left ( \\frac{\\nu_{i2}^0}{2}, \\frac{S_{i2}^0}{2} \\right )\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "where $\\mathcal{IW}$ denotes the inverse-Wishart distribution and $\\mathcal{IG}$ denotes the inverse-Gamma distribution. We set the prior hyperparameters as:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "v_1^0 = T + 3, & \\quad S_1^0 = I \\\\\n",
    "v_{i2}^0 = 6, & \\quad S_{i2}^0 = 0.01 \\qquad \\text{for each} ~ i\\\\\n",
    "\\end{aligned}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prior hyperparameters\n",
    "\n",
    "# Prior for obs. cov. is inverse-Wishart(v_1^0=k + 3, S10=I)\n",
    "v10 = mod.k_endog + 3\n",
    "S10 = np.eye(mod.k_endog)\n",
    "\n",
    "# Prior for state cov. variances is inverse-Gamma(v_{i2}^0 / 2 = 3, S+{i2}^0 / 2 = 0.005)\n",
    "vi20 = 6\n",
    "Si20 = 0.01"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before running the MCMC iterations, there are a couple of practical steps:\n",
    "\n",
    "1. Create arrays to store the draws of our state vector, observation covariance matrix, and state error variances.\n",
    "2. Put the initial values for H and Q (described above) into the storage vectors\n",
    "3. Construct the simulation smoother object associated with our `TVPVAR` instance to make draws of the state vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Gibbs sampler setup\n",
    "niter = 11000\n",
    "nburn = 1000\n",
    "\n",
    "# 1. Create storage arrays\n",
    "store_states = np.zeros((niter + 1, mod.nobs, mod.k_states))\n",
    "store_obs_cov = np.zeros((niter + 1, mod.k_endog, mod.k_endog))\n",
    "store_state_cov = np.zeros((niter + 1, mod.k_states))\n",
    "\n",
    "# 2. Put in the initial values\n",
    "store_obs_cov[0] = initial_obs_cov\n",
    "store_state_cov[0] = initial_state_cov_diag\n",
    "mod.update_variances(store_obs_cov[0], store_state_cov[0])\n",
    "\n",
    "# 3. Construct posterior samplers\n",
    "sim = mod.simulation_smoother(method='cfa')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As before, we could have used either the simulation smoother based on the Kalman filter and smoother or that based on the Cholesky Factor Algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(niter):\n",
    "    mod.update_variances(store_obs_cov[i], store_state_cov[i])\n",
    "    sim.simulate()\n",
    "\n",
    "    # 1. Sample states\n",
    "    store_states[i + 1] = sim.simulated_state.T\n",
    "\n",
    "    # 2. Simulate obs cov\n",
    "    fitted = np.matmul(mod['design'].transpose(2, 0, 1), store_states[i + 1][..., None])[..., 0]\n",
    "    resid = mod.endog - fitted\n",
    "    store_obs_cov[i + 1] = invwishart.rvs(v10 + mod.nobs, S10 + resid.T @ resid)\n",
    "\n",
    "    # 3. Simulate state cov variances\n",
    "    resid = store_states[i + 1, 1:] - store_states[i + 1, :-1]\n",
    "    sse = np.sum(resid**2, axis=0)\n",
    "    \n",
    "    for j in range(mod.k_states):\n",
    "        rv = invgamma.rvs((vi20 + mod.nobs - 1) / 2, scale=(Si20 + sse[j]) / 2)\n",
    "        store_state_cov[i + 1, j] = rv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After removing a number of initial draws, the remaining draws from the posterior allow us to conduct inference. Below, we plot the posterior mean of the time-varying regression coefficients.\n",
    "\n",
    "(**Note**: these plots are different from those in Figure 1 of the published version of Chan and Jeliakov (2009), but they are very similar to those produced by the Matlab replication code available at http://joshuachan.org/code/code_TVPVAR.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Collect the posterior means of each time-varying coefficient\n",
    "states_posterior_mean = pd.DataFrame(\n",
    "    np.mean(store_states[nburn + 1:], axis=0),\n",
    "    index=mod._index, columns=mod.state_names)\n",
    "\n",
    "# Plot these means over time\n",
    "plot_coefficients_by_equation(states_posterior_mean);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Python also has a number of libraries to assist with exploring Bayesian models. Here we'll just use the [arviz](https://arviz-devs.github.io/arviz/index.html) package to explore the credible intervals of each of the covariance and variance parameters, although it makes available a much wider set of tools for analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import arviz as az\n",
    "\n",
    "# Collect the observation error covariance parameters\n",
    "az_obs_cov = az.convert_to_inference_data({\n",
    "    ('Var[%s]' % mod.endog_names[i] if i == j else\n",
    "     'Cov[%s, %s]' % (mod.endog_names[i], mod.endog_names[j])):\n",
    "    store_obs_cov[nburn + 1:, i, j]\n",
    "    for i in range(mod.k_endog) for j in range(i, mod.k_endog)})\n",
    "\n",
    "# Plot the credible intervals\n",
    "az.plot_forest(az_obs_cov, figsize=(8, 7));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Collect the state innovation variance parameters\n",
    "az_state_cov = az.convert_to_inference_data({\n",
    "    r'$\\sigma^2$[%s]' % mod.state_names[i]: store_state_cov[nburn + 1:, i]\n",
    "    for i in range(mod.k_states)})\n",
    "\n",
    "# Plot the credible intervals\n",
    "az.plot_forest(az_state_cov, figsize=(8, 7));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Appendix: performance\n",
    "\n",
    "Finally, we run a few simple tests to compare the performance of the KFS and CFA simulation smoothers by using the `%timeit` Jupytor notebook magic.\n",
    "\n",
    "One caveat is that the KFS simulation smoother can produce a variety of output beyond just simulations of the posterior state vector, and these additional computations could bias the results. To make the results comparable, we will tell the KFS simulation smoother to only compute simulations of the state by using the `simulation_output` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from statsmodels.tsa.statespace.simulation_smoother import SIMULATION_STATE\n",
    "\n",
    "sim_cfa = mod.simulation_smoother(method='cfa')\n",
    "sim_kfs = mod.simulation_smoother(simulation_output=SIMULATION_STATE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we can use the following code to perform a basic timing exercise:\n",
    "\n",
    "```python\n",
    "%timeit -n 10000 -r 3 sim_cfa.simulate()\n",
    "%timeit -n 10000 -r 3 sim_kfs.simulate()\n",
    "```\n",
    "\n",
    "On the machine this was tested on, this resulted in the following:\n",
    "\n",
    "```\n",
    "2.06 ms ± 26.5 µs per loop (mean ± std. dev. of 3 runs, 10000 loops each)\n",
    "2.02 ms ± 68.4 µs per loop (mean ± std. dev. of 3 runs, 10000 loops each)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These results suggest that - at least for this model - there are not noticeable computational gains from the CFA approach relative to the KFS approach. However, this does not rule out the following:\n",
    "\n",
    "1. The Statsmodels implementation of the CFA simulation smoother could possibly be further optimized\n",
    "2. The CFA apporach may only show improvement for certain models (for example with a large number of `endog` variables)\n",
    "\n",
    "One simple way to take a first pass at assessing the first possibility is to compare the runtime of the Statsmodels implementation of the CFA simulation smoother to the Matlab implementation in the replication codes of Chan and Jeliakov (2009), available at http://joshuachan.org/code/code_TVPVAR.html.\n",
    "\n",
    "While the Statsmodels version of the CFA simulation smoother is written in Cython and compiled to C code, the Matlab version takes advantage of the Matlab's sparse matrix capabilities. As a result, even though it is not compiled code,  we might expect it to have relatively good performance.\n",
    "\n",
    "On the machine this was tested on, the Matlab version typically ran the MCMC loop with 11,000 iterations in 70-75 seconds, while the MCMC loop in this notebook using the Statsmodels CFA simulation smoother (see above), also with 11,0000 iterations, ran in 40-45 seconds. This is some evidence that the Statsmodels implementation of the CFA smoother already performs relatively well (although it doesn't rule out that there are additional gains possible)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bibliography\n",
    "\n",
    "Carter, Chris K., and Robert Kohn. \"On Gibbs sampling for state space models.\" Biometrika 81, no. 3 (1994): 541-553.\n",
    "\n",
    "Chan, Joshua CC, and Ivan Jeliazkov. \"Efficient simulation and integrated likelihood estimation in state space models.\" International Journal of Mathematical Modelling and Numerical Optimisation 1, no. 1-2 (2009): 101-120.\n",
    "\n",
    "De Jong, Piet, and Neil Shephard. \"The simulation smoother for time series models.\" Biometrika 82, no. 2 (1995): 339-350.\n",
    "\n",
    "Durbin, James, and Siem Jan Koopman. \"A simple and efficient simulation smoother for state space time series analysis.\" Biometrika 89, no. 3 (2002): 603-616.\n",
    "\n",
    "McCausland, William J., Shirley Miller, and Denis Pelletier. \"Simulation smoothing for state–space models: A computational efficiency analysis.\" Computational Statistics & Data Analysis 55, no. 1 (2011): 199-212."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
